use once_cell::sync::Lazy;
use std::collections::HashMap;
use std::fmt;
use std::str::Chars;
use std::sync::RwLock;

#[derive(Clone, Debug)]
pub enum Expression {
    Constant(f64),
    Variable(Variable),
    Unary(UnaryOperator, Box<Expression>),
    Binary(BinaryOperator, Box<Expression>, Box<Expression>),
}

#[derive(Clone, Copy, Debug)]
pub enum Variable {
    X,
    Y,
}

#[derive(Clone, Copy, Debug)]
pub enum UnaryOperator {
    Negate,
    Function(UnaryFunction),
}

#[derive(Clone, Copy, Debug)]
pub enum UnaryFunction {
    Sin,
    Cos,
    Tan,
    Asin,
    Acos,
    Atan,
    Exp,
    Abs,
    Ln,
    Log10,
    Sqrt,
}

#[derive(Clone, Copy, Debug)]
pub enum BinaryOperator {
    Add,
    Sub,
    Mul,
    Div,
    Pow,
}

#[derive(Debug)]
pub enum CoreError {
    Parse(ParseError),
    Evaluate(EvaluateError),
}

impl fmt::Display for CoreError {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            CoreError::Parse(err) => write!(f, "parse error: {}", err),
            CoreError::Evaluate(err) => write!(f, "evaluation error: {}", err),
        }
    }
}

#[derive(Debug)]
pub enum ParseError {
    UnexpectedEnd,
    UnexpectedToken(String),
    InvalidNumber(String),
    UnknownIdentifier(String),
    MismatchedParenthesis,
    ArgumentCount {
        name: String,
        expected: usize,
        got: usize,
    },
}

impl fmt::Display for ParseError {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            ParseError::UnexpectedEnd => write!(f, "expression ended unexpectedly"),
            ParseError::UnexpectedToken(token) => write!(f, "unexpected token `{}`", token),
            ParseError::InvalidNumber(num) => write!(f, "invalid number literal `{}`", num),
            ParseError::UnknownIdentifier(ident) => write!(f, "unknown identifier `{}`", ident),
            ParseError::MismatchedParenthesis => write!(f, "unmatched parenthesis"),
            ParseError::ArgumentCount {
                name,
                expected,
                got,
            } => {
                write!(
                    f,
                    "{} expected {} argument(s) but received {}",
                    name, expected, got
                )
            }
        }
    }
}

#[derive(Debug)]
pub enum EvaluateError {
    DomainError(&'static str),
    Indeterminate,
}

impl fmt::Display for EvaluateError {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            EvaluateError::DomainError(msg) => write!(f, "{}", msg),
            EvaluateError::Indeterminate => write!(f, "result is NaN or infinite"),
        }
    }
}

struct Lexer<'a> {
    input: &'a str,
    chars: Chars<'a>,
    lookahead: Option<char>,
    index: usize,
}

impl<'a> Lexer<'a> {
    fn new(input: &'a str) -> Self {
        Self {
            input,
            chars: input.chars(),
            lookahead: None,
            index: 0,
        }
    }

    fn peek_char(&mut self) -> Option<char> {
        if let Some(ch) = self.lookahead {
            Some(ch)
        } else {
            self.lookahead = self.chars.next();
            self.lookahead
        }
    }

    fn next_char(&mut self) -> Option<char> {
        if let Some(ch) = self.lookahead.take() {
            self.index += ch.len_utf8();
            Some(ch)
        } else {
            let ch = self.chars.next()?;
            self.index += ch.len_utf8();
            Some(ch)
        }
    }

    fn skip_whitespace(&mut self) {
        while matches!(self.peek_char(), Some(ch) if ch.is_whitespace()) {
            self.next_char();
        }
    }

    fn lex(&mut self) -> Result<Vec<Token>, ParseError> {
        let mut tokens = Vec::new();
        while let Some(ch) = self.peek_char() {
            if ch.is_whitespace() {
                self.skip_whitespace();
                continue;
            }
            match ch {
                '0'..='9' | '.' => tokens.push(self.lex_number()?),
                'a'..='z' | 'A'..='Z' | '_' => tokens.push(self.lex_identifier()),
                '+' => {
                    self.next_char();
                    tokens.push(Token::Plus);
                }
                '-' => {
                    self.next_char();
                    tokens.push(Token::Minus);
                }
                '*' => {
                    self.next_char();
                    tokens.push(Token::Star);
                }
                '/' => {
                    self.next_char();
                    tokens.push(Token::Slash);
                }
                '^' => {
                    self.next_char();
                    tokens.push(Token::Caret);
                }
                '(' => {
                    self.next_char();
                    tokens.push(Token::LeftParen);
                }
                ')' => {
                    self.next_char();
                    tokens.push(Token::RightParen);
                }
                ',' => {
                    self.next_char();
                    tokens.push(Token::Comma);
                }
                _ => {
                    return Err(ParseError::UnexpectedToken(ch.to_string()));
                }
            }
        }
        Ok(tokens)
    }

    fn lex_number(&mut self) -> Result<Token, ParseError> {
        let start_index = self.index;
        let mut literal = String::new();
        let mut has_decimal_point = false;

        if let Some('.') = self.peek_char() {
            has_decimal_point = true;
            literal.push('.');
            self.next_char();
        }

        while let Some(ch) = self.peek_char() {
            if ch.is_ascii_digit() {
                literal.push(ch);
                self.next_char();
            } else if ch == '.' && !has_decimal_point {
                has_decimal_point = true;
                literal.push('.');
                self.next_char();
            } else {
                break;
            }
        }

        if let Some(ch) = self.peek_char() {
            if matches!(ch, 'e' | 'E') {
                literal.push(ch);
                self.next_char();
                if let Some(sign @ ('+' | '-')) = self.peek_char() {
                    literal.push(sign);
                    self.next_char();
                }
                let mut has_exponent_digits = false;
                while let Some(exp_ch) = self.peek_char() {
                    if exp_ch.is_ascii_digit() {
                        has_exponent_digits = true;
                        literal.push(exp_ch);
                        self.next_char();
                    } else {
                        break;
                    }
                }
                if !has_exponent_digits {
                    let fragment = &self.input[start_index..self.index];
                    return Err(ParseError::InvalidNumber(fragment.to_string()));
                }
            }
        }

        literal
            .parse::<f64>()
            .map(Token::Number)
            .map_err(|_| ParseError::InvalidNumber(literal))
    }

    fn lex_identifier(&mut self) -> Token {
        let mut ident = String::new();
        while let Some(ch) = self.peek_char() {
            if ch.is_ascii_alphanumeric() || ch == '_' {
                ident.push(ch);
                self.next_char();
            } else {
                break;
            }
        }
        Token::Identifier(ident)
    }
}

#[derive(Clone, Debug)]
enum Token {
    Number(f64),
    Identifier(String),
    Plus,
    Minus,
    Star,
    Slash,
    Caret,
    LeftParen,
    RightParen,
    Comma,
}

struct TokenStream {
    tokens: Vec<Token>,
    index: usize,
}

impl TokenStream {
    fn new(tokens: Vec<Token>) -> Self {
        Self { tokens, index: 0 }
    }

    fn peek(&self) -> Option<&Token> {
        self.tokens.get(self.index)
    }

    fn next(&mut self) -> Option<Token> {
        let token = self.tokens.get(self.index).cloned();
        if token.is_some() {
            self.index += 1;
        }
        token
    }

    fn expect_right_paren(&mut self) -> Result<(), ParseError> {
        match self.next() {
            Some(Token::RightParen) => Ok(()),
            _ => Err(ParseError::MismatchedParenthesis),
        }
    }
}

static EXPRESSION_CACHE: Lazy<RwLock<HashMap<String, Expression>>> =
    Lazy::new(|| RwLock::new(HashMap::new()));

pub fn evaluate(expression: &str, x: f64, y: f64) -> Result<f64, CoreError> {
    let expression = expression.trim();
    if expression.is_empty() {
        return Err(CoreError::Parse(ParseError::UnexpectedEnd));
    }

    let ast = {
        let cache = EXPRESSION_CACHE.read().unwrap();
        if let Some(existing) = cache.get(expression) {
            existing.clone()
        } else {
            drop(cache);
            let parsed = parse_expression(expression)?;
            let mut cache = EXPRESSION_CACHE.write().unwrap();
            cache.insert(expression.to_string(), parsed.clone());
            parsed
        }
    };

    let value = evaluate_ast(&ast, x, y)?;
    if value.is_finite() {
        Ok(value)
    } else {
        Err(CoreError::Evaluate(EvaluateError::Indeterminate))
    }
}

fn parse_expression(input: &str) -> Result<Expression, CoreError> {
    if let Some((left, right)) = split_equation(input) {
        let left_expr = parse_expression(&left)?;
        let right_expr = parse_expression(&right)?;
        return Ok(Expression::Binary(
            BinaryOperator::Sub,
            Box::new(left_expr),
            Box::new(right_expr),
        ));
    }

    let mut lexer = Lexer::new(input);
    let tokens = lexer.lex().map_err(CoreError::Parse)?;
    let mut stream = TokenStream::new(tokens);
    let ast = parse_bp(&mut stream, 0).map_err(CoreError::Parse)?;
    if stream.peek().is_some() {
        return Err(CoreError::Parse(ParseError::UnexpectedToken(format!(
            "{:?}",
            stream.peek().unwrap()
        ))));
    }
    Ok(ast)
}

fn parse_bp(stream: &mut TokenStream, min_bp: u8) -> Result<Expression, ParseError> {
    let mut lhs = parse_prefix(stream)?;
    loop {
        let token = match stream.peek() {
            Some(token) => token.clone(),
            None => break,
        };
        let (l_bp, r_bp, op) = match infix_binding_power(&token) {
            Some(info) => info,
            None => break,
        };
        if l_bp < min_bp {
            break;
        }
        stream.next();
        let rhs = parse_bp(stream, r_bp)?;
        lhs = Expression::Binary(op, Box::new(lhs), Box::new(rhs));
    }
    Ok(lhs)
}

fn parse_prefix(stream: &mut TokenStream) -> Result<Expression, ParseError> {
    match stream.next() {
        Some(Token::Number(value)) => Ok(Expression::Constant(value)),
        Some(Token::Identifier(name)) => parse_identifier(name, stream),
        Some(Token::Minus) => {
            let expr = parse_bp(stream, 6)?;
            Ok(Expression::Unary(UnaryOperator::Negate, Box::new(expr)))
        }
        Some(Token::Plus) => parse_bp(stream, 6),
        Some(Token::LeftParen) => {
            let expr = parse_bp(stream, 0)?;
            stream.expect_right_paren()?;
            Ok(expr)
        }
        Some(other) => Err(ParseError::UnexpectedToken(format!("{:?}", other))),
        None => Err(ParseError::UnexpectedEnd),
    }
}

fn parse_identifier(name: String, stream: &mut TokenStream) -> Result<Expression, ParseError> {
    let name_lower = name.to_lowercase();
    let is_function_call = matches!(stream.peek(), Some(Token::LeftParen));

    if is_function_call {
        stream.next(); // consume '('
        let args = parse_arguments(stream)?;
        return build_function_call(&name_lower, args);
    }

    match name_lower.as_str() {
        "x" => Ok(Expression::Variable(Variable::X)),
        "y" => Ok(Expression::Variable(Variable::Y)),
        "pi" => Ok(Expression::Constant(std::f64::consts::PI)),
        "e" => Ok(Expression::Constant(std::f64::consts::E)),
        _ => Err(ParseError::UnknownIdentifier(name)),
    }
}

fn parse_arguments(stream: &mut TokenStream) -> Result<Vec<Expression>, ParseError> {
    let mut args = Vec::new();
    if matches!(stream.peek(), Some(Token::RightParen)) {
        stream.next();
        return Ok(args);
    }

    loop {
        let expr = parse_bp(stream, 0)?;
        args.push(expr);
        match stream.peek() {
            Some(Token::Comma) => {
                stream.next();
                continue;
            }
            Some(Token::RightParen) => {
                stream.next();
                break;
            }
            _ => return Err(ParseError::MismatchedParenthesis),
        }
    }
    Ok(args)
}

fn build_function_call(name: &str, args: Vec<Expression>) -> Result<Expression, ParseError> {
    match name {
        "sin" => expect_unary_call(name, args, UnaryFunction::Sin),
        "cos" => expect_unary_call(name, args, UnaryFunction::Cos),
        "tan" => expect_unary_call(name, args, UnaryFunction::Tan),
        "asin" => expect_unary_call(name, args, UnaryFunction::Asin),
        "acos" => expect_unary_call(name, args, UnaryFunction::Acos),
        "atan" => expect_unary_call(name, args, UnaryFunction::Atan),
        "exp" => expect_unary_call(name, args, UnaryFunction::Exp),
        "abs" => expect_unary_call(name, args, UnaryFunction::Abs),
        "ln" => expect_unary_call(name, args, UnaryFunction::Ln),
        "log" => expect_unary_call(name, args, UnaryFunction::Log10),
        "sqrt" => expect_unary_call(name, args, UnaryFunction::Sqrt),
        "pow" => expect_binary_call(name, args, BinaryOperator::Pow),
        _ => Err(ParseError::UnknownIdentifier(name.to_string())),
    }
}

fn expect_unary_call(
    name: &str,
    args: Vec<Expression>,
    func: UnaryFunction,
) -> Result<Expression, ParseError> {
    if args.len() != 1 {
        return Err(ParseError::ArgumentCount {
            name: name.to_string(),
            expected: 1,
            got: args.len(),
        });
    }
    let arg = args.into_iter().next().unwrap();
    Ok(Expression::Unary(
        UnaryOperator::Function(func),
        Box::new(arg),
    ))
}

fn expect_binary_call(
    name: &str,
    mut args: Vec<Expression>,
    op: BinaryOperator,
) -> Result<Expression, ParseError> {
    if args.len() != 2 {
        return Err(ParseError::ArgumentCount {
            name: name.to_string(),
            expected: 2,
            got: args.len(),
        });
    }
    let right = args.pop().unwrap();
    let left = args.pop().unwrap();
    Ok(Expression::Binary(op, Box::new(left), Box::new(right)))
}

fn infix_binding_power(token: &Token) -> Option<(u8, u8, BinaryOperator)> {
    match token {
        Token::Plus => Some((1, 2, BinaryOperator::Add)),
        Token::Minus => Some((1, 2, BinaryOperator::Sub)),
        Token::Star => Some((3, 4, BinaryOperator::Mul)),
        Token::Slash => Some((3, 4, BinaryOperator::Div)),
        Token::Caret => Some((6, 5, BinaryOperator::Pow)),
        _ => None,
    }
}

fn split_equation(expression: &str) -> Option<(String, String)> {
    let mut depth = 0usize;
    for (idx, ch) in expression.char_indices() {
        match ch {
            '(' => depth += 1,
            ')' => depth = depth.saturating_sub(1),
            '=' if depth == 0 => {
                let left = expression[..idx].trim().to_string();
                let right = expression[idx + ch.len_utf8()..].trim().to_string();
                if !left.is_empty() && !right.is_empty() {
                    return Some((left, right));
                }
            }
            _ => {}
        }
    }
    None
}

fn evaluate_ast(expr: &Expression, x: f64, y: f64) -> Result<f64, CoreError> {
    match expr {
        Expression::Constant(value) => Ok(*value),
        Expression::Variable(Variable::X) => Ok(x),
        Expression::Variable(Variable::Y) => Ok(y),
        Expression::Unary(operator, inner) => {
            let value = evaluate_ast(inner, x, y)?;
            let result = match operator {
                UnaryOperator::Negate => -value,
                UnaryOperator::Function(func) => apply_unary_function(*func, value)?,
            };
            Ok(result)
        }
        Expression::Binary(operator, left, right) => {
            let lhs = evaluate_ast(left, x, y)?;
            let rhs = evaluate_ast(right, x, y)?;
            let result = match operator {
                BinaryOperator::Add => lhs + rhs,
                BinaryOperator::Sub => lhs - rhs,
                BinaryOperator::Mul => lhs * rhs,
                BinaryOperator::Div => {
                    if rhs == 0.0 {
                        return Err(CoreError::Evaluate(EvaluateError::DomainError(
                            "division by zero is undefined",
                        )));
                    }
                    lhs / rhs
                }
                BinaryOperator::Pow => lhs.powf(rhs),
            };
            if result.is_finite() {
                Ok(result)
            } else {
                Err(CoreError::Evaluate(EvaluateError::Indeterminate))
            }
        }
    }
}

fn apply_unary_function(func: UnaryFunction, value: f64) -> Result<f64, CoreError> {
    use std::f64::consts::LN_10;
    let result = match func {
        UnaryFunction::Sin => value.sin(),
        UnaryFunction::Cos => value.cos(),
        UnaryFunction::Tan => value.tan(),
        UnaryFunction::Asin => {
            if !(-1.0..=1.0).contains(&value) {
                return Err(CoreError::Evaluate(EvaluateError::DomainError(
                    "asin expects argument in [-1, 1]",
                )));
            }
            value.asin()
        }
        UnaryFunction::Acos => {
            if !(-1.0..=1.0).contains(&value) {
                return Err(CoreError::Evaluate(EvaluateError::DomainError(
                    "acos expects argument in [-1, 1]",
                )));
            }
            value.acos()
        }
        UnaryFunction::Atan => value.atan(),
        UnaryFunction::Exp => value.exp(),
        UnaryFunction::Abs => value.abs(),
        UnaryFunction::Ln => {
            if value <= 0.0 {
                return Err(CoreError::Evaluate(EvaluateError::DomainError(
                    "ln expects argument > 0",
                )));
            }
            value.ln()
        }
        UnaryFunction::Log10 => {
            if value <= 0.0 {
                return Err(CoreError::Evaluate(EvaluateError::DomainError(
                    "log expects argument > 0",
                )));
            }
            value.ln() / LN_10
        }
        UnaryFunction::Sqrt => {
            if value < 0.0 {
                return Err(CoreError::Evaluate(EvaluateError::DomainError(
                    "sqrt expects argument >= 0",
                )));
            }
            value.sqrt()
        }
    };
    if result.is_finite() {
        Ok(result)
    } else {
        Err(CoreError::Evaluate(EvaluateError::Indeterminate))
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    fn eval(expr: &str, x: f64, y: f64) -> f64 {
        evaluate(expr, x, y).expect("expression to evaluate")
    }

    #[test]
    fn evaluates_basic_arithmetic() {
        assert_eq!(eval("1 + 2 * 3", 0.0, 0.0), 7.0);
        assert_eq!(eval("(1 + 2) * 3", 0.0, 0.0), 9.0);
        assert!((eval("2^3", 0.0, 0.0) - 8.0).abs() < 1e-6);
    }

    #[test]
    fn evaluates_variables_and_equations() {
        assert_eq!(eval("x + y", 1.0, 2.0), 3.0);
        assert_eq!(eval("x^2 + y^2 = 1", 1.0, 0.0), 0.0);
    }

    #[test]
    fn evaluates_functions() {
        let result = eval("sin(pi / 2)", 0.0, 0.0);
        assert!((result - 1.0).abs() < 1e-6);
        let ln_exp = eval("ln(exp(3))", 0.0, 0.0);
        assert!((ln_exp - 3.0).abs() < 1e-6);
        let pow_form = eval("pow(2, 4)", 0.0, 0.0);
        assert!((pow_form - 16.0).abs() < 1e-6);
    }

    #[test]
    fn caches_parsed_expression() {
        let first = eval("x + y", 2.0, 3.0);
        let second = eval("x + y", 4.0, 5.0);
        assert_eq!(first, 5.0);
        assert_eq!(second, 9.0);
        let cache = EXPRESSION_CACHE.read().unwrap();
        assert!(cache.contains_key("x + y"));
    }
}
